{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copyright 2019 Google LLC\n",
    "# \n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "#     https://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a target=\"_blank\" href=\"https://colab.research.google.com/github/GoogleCloudPlatform/keras-idiomatic-programmer/blob/master/notebooks/prediction_with_example_key.ipynb\">\n",
    "<img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Batch Prediction - Adding Identifier Key to each Example\n",
    "\n",
    "Recently, someone asked me a problem related to doing distributed prediction on a batch prediction request, and returning the results in real-time. The problem they had is that the backend of their application which sends the batch request does not know what order the individual predictions will arrive in, and thus which prediction result when with which prediction request.\n",
    "\n",
    "In this case, they had an existing trained tf.keras model and wanted to retrofit the model such that each individual prediction request had a unique identifier and the prediction result includes the corresponding unique identifier; which we will refer herein as the example key.\n",
    "\n",
    "This is so easy in Keras using the Functional API with multiple inputs and multiple outputs, and using the trained Keras model object as a callable (i.e., a layer). Here's the basic steps we will follow:\n",
    "\n",
    "    1. Build and train the existing model.\n",
    "    2. Make a second model (we call the wrapper model), with\n",
    "        A. Two inputs\n",
    "        B. Two outputs\n",
    "    3. The two inputs to the wrapper model are:\n",
    "        A. The input layer of the trained model\n",
    "        B. A 1D vector for the example key\n",
    "    4. The two outputs to the wrapper model are:\n",
    "        A. The output from the trained model\n",
    "        B. The value (example key) of the 1D vector.\n",
    "        \n",
    "That's all you need. Let's show an example of doing this with a tf.Keras model for MNIST.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "We start by importing the tf.keras modules we will use, along with numpy and the builtin dataset for MNIST.\n",
    "\n",
    "This tutorial will work with both TF 1.X and TF 2.0."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "from tensorflow.keras import Input, Model\n",
    "from tensorflow.keras.layers import Dense, Flatten\n",
    "from tensorflow.keras.datasets import mnist\n",
    "from tensorflow.keras.utils import to_categorical\n",
    "import numpy as np\n",
    "\n",
    "# Reguired if you are using pre TF-2.0\n",
    "# i.e, eager execution is required for adding example key as a graph operation\n",
    "tf.enable_eager_execution()\n",
    "\n",
    "# Display what version of TF you are using\n",
    "print(tf.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare Data\n",
    "\n",
    "Next, we get the data from the tf.keras builtin dataset for MNIST and do the standard preprocessing:\n",
    "\n",
    "    1. Normalize the image data (divide by 255) and set datatype to FP32\n",
    "    2. One hot encode the labels (to_categorical)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(x_train, y_train), (x_test, y_test) = mnist.load_data()\n",
    "x_train = (x_train / 255.0).astype(np.float32)\n",
    "x_test  = (x_test  / 255.0).astype(np.float32)\n",
    "y_train = to_categorical(y_train)\n",
    "y_test  = to_categorical(y_test)\n",
    "\n",
    "print(\"Train\", x_train.shape, y_train.shape)\n",
    "\n",
    "print(\"Test \", x_test.shape,  y_test.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build MNIST Model\n",
    "\n",
    "We will use the Functional API to build essential a sequential dense neural network (DNN), as follows:\n",
    "\n",
    "    1. Flattens the MNIST images (28x28) into a 1D vector (784).\n",
    "    2. An input and hidden dense layer of 128 nodes with a ReLU activation function.\n",
    "    3. An output layer of 10 nodes (1 per digit) with a softmax activation function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = Input((28, 28))\n",
    "\n",
    "x = Flatten()(inputs)\n",
    "x = Dense(128, activation='relu')(x)\n",
    "x = Dense(128, activation='relu')(x)\n",
    "outputs = Dense(10, activation='softmax')(x)\n",
    "\n",
    "model = Model(inputs, outputs)\n",
    "model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['acc'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train the MNIST Model\n",
    "\n",
    "Let's do a quick training of this MNIST model for 5 epochs. We should see ~98% accuracy on the training data. Let's not worry about the test data or overfitting. We just need a quick model to demonstrate how to make a wrapper model using the tf.Keras Functional API and pass through the example key during prediction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Train the model\n",
    "model.fit(x_train, y_train, epochs=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Display the Model Architecture\n",
    "\n",
    "We will use the summary() method to display the model architecture."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Functional API and multiple inputs / outputs\n",
    "\n",
    "If you have not used multiple inputs and multiple outputs, this maybe the easiest example you will see. In the Functional API, when we build the model using the Model() class, we pass two parameters, the input tensor and the output layer; which I call pulling it all together connecting the inputs to the outputs:\n",
    "\n",
    "    my_model = Model(inputs, outputs)\n",
    "    \n",
    "I won't go into detail about how to (and why) to create models with more than one input and/or output. We will simply use this method to implement passing through an example key at prediction with an existing trained tf.keras model (i.e., MNIST model). The syntax for specifying both multiple inputs and outputs looks like this:\n",
    "\n",
    "    my_model = Model( [inputs1, inputs2], [outputs1, outputs2])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build the Wrapper Model\n",
    "\n",
    "Let's get started. We can do this in two lines of Keras code!\n",
    "\n",
    "1) First, we need to define a second input for passing through the example key. It will need to be number. For prediction purposes, we will define this as a scalar value. We will name this input tensor as *key*.\n",
    "\n",
    "            key = Input(1)\n",
    "            \n",
    "2) Next. we build a new model (\"the wrapper\") reusing the existing trained model, but we are going to add the key tensor to both the inputs and outputs. To resuse the trained model, we need to get the trained model's input tensor and output layer, as in:\n",
    "\n",
    "            Model(model.input, model.output)\n",
    "            \n",
    "Let's now expand on this and add in the key:\n",
    "\n",
    "            wrapper_model = Model( [model.input, key], (model.output, key])\n",
    "            \n",
    "That's it. We did it. We don't need to compile or train --because we reused the tf.keras model that is already trained!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the second input to the model for passing through the example key to the output\n",
    "key = Input((1))\n",
    "\n",
    "# Build the wrapper model\n",
    "wrapper_model = Model([model.input, key], [model.output, key])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Make a Prediction\n",
    "\n",
    "Let's now make a prediction using the wrapper model. We will pass an arbitrary example from the test data along with a unique ID to the wrapper model's *predict()* method.\n",
    "\n",
    "            result = wrapper_model.predict( [example, key] )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's grap an arbitrary example from the test data.\n",
    "example = x_test[7]\n",
    "# Eventhough it is one example, we need to fit it into a batch (of one)\n",
    "example = np.expand_dims(example, axis=0)\n",
    "print(example.shape)\n",
    "\n",
    "# Let's make an 1D input for our key and choose a unique value --in this case, we use the number 45\n",
    "example_key = np.asarray([45])\n",
    "print(example_key.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's make the call now. We will get back a tuple. The first value is the prediction (10 element array --one for each digit) and the example key which is passed through (i.e., 45)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Do the prediction\n",
    "prediction, id = wrapper_model.predict([example, example_key])\n",
    "\n",
    "print(prediction)\n",
    "\n",
    "print(\"The Example Key\", id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We did it. You can see the value 45 in the returned id :)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
